/*BEGIN_LEGAL
 *
 * There should be some kind of license here.
 * ===================================================================== */
/*
  @ORIGINAL_AUTHOR: Benjamin Ylvisaker
*/

#include <errno.h>
#include "pin.H"
extern "C" {
#include "xed-interface.h"
}
#include <list>
#include <iostream>
#include <fstream>
#include <stdlib.h>
#include <stdint.h>
#include <string.h>

/* Generic utilities */

#ifdef __GNUC__
#define VARIABLE_IS_NOT_USED __attribute__ ((unused))
#else
#define VARIABLE_IS_NOT_USED
#endif

/* "ff_" fail-fast versions of memory allocation functions */

static inline void *ff_malloc( size_t s )
{
    void *p = malloc( s );
    if( !p )
    {
        PIN_ExitProcess( ENOMEM );
    }
    return p;
}

static inline void *ff_calloc( size_t n, size_t s )
{
    void *p = calloc( n, s );
    if( !p )
    {
        PIN_ExitProcess( ENOMEM );
    }
    return p;
}

static inline void *ff_realloc( void *r, size_t s )
{
    void *p = realloc( r, s );
    if( !p )
    {
        std::cout << "realloc failed! ptr:" << r << "  sz:" << s;
        PIN_ExitProcess( ENOMEM );
    }
    return p;
}

#define CHECK_CAP_REALLOC( p, t, i, c ) \
    do { \
        if( i >= c ) \
        { \
            uint32_t old_cap = c; \
            c *= 2; \
            p = (t)ff_realloc( p, c * sizeof( p[0] ) );\
            memset( &p[ old_cap ], 0, old_cap * sizeof( p[0] ) ); \
        } \
    } while( 0 )

/* ===================================================================== */
/* Commandline Switches */
/* ===================================================================== */

KNOB<string> KnobOutputLogName(
    KNOB_MODE_WRITEONCE,
    "pintool",
    "logname",
    "bens_miraculous_log",
    "specify log name" );
KNOB<string> KnobOutputFileBase(
    KNOB_MODE_WRITEONCE,
    "pintool",
    "o_base",
    "bbl_cfg",
    "specify profile file name base" );
KNOB<string> KnobOutputFileSuffix(
    KNOB_MODE_WRITEONCE,
    "pintool",
    "o_suffix",
    "json",
    "specify profile file name suffix" );

/* ===================================================================== */
/* Print Help Message                                                    */
/* ===================================================================== */

INT32 Usage()
{
    cerr <<
        "This pin tool is Ben's queue nonsense\n"
        "\n";

    cerr << KNOB_BASE::StringKnobSummary();

    cerr << endl;

    return -1;
}

/* ===================================================================== */
/* Global Variables */
/* ===================================================================== */

typedef struct _EDGE _EDGE, *EDGE;
struct _EDGE
{
    uint32_t id, from, to;
    uint64_t exec_ct;
};

typedef uint32_t AREG_FLAGS;
const AREG_FLAGS REGF_NONE       =  0x0;
const AREG_FLAGS REGF_READ_WRITE =  0x1;
const AREG_FLAGS REGF_IMPLICIT   =  0x2;

typedef struct _ANNOT_REG _ANNOT_REG, *ANNOT_REG;
struct _ANNOT_REG
{
    REG r;
    AREG_FLAGS flags;
};

typedef uint32_t INSTR_FLAGS;
const INSTR_FLAGS INSTRF_NONE                    =  0x00;
const INSTR_FLAGS INSTRF_HAS_FALL                =  0x01;
const INSTR_FLAGS INSTRF_MEM_READ                =  0x02;
const INSTR_FLAGS INSTRF_MEM_WRITE               =  0x04;
const INSTR_FLAGS INSTRF_TOO_MANY_READ_OPERANDS  =  0x08;
const INSTR_FLAGS INSTRF_TOO_MANY_WRITE_OPERANDS =  0x10;

typedef struct _INSTR _INSTR, *INSTR;
struct _INSTR
{
    char *disas;
    OPCODE opc;
    uint8_t r_ct, w_ct;
    ANNOT_REG reads, writes;
    INSTR_FLAGS flags;
};

typedef uint32_t BBL_FLAGS;
const BBL_FLAGS BBLF_NONE           =  0x0;
const BBL_FLAGS BBLF_IS_CALL        =  0x1;
const BBL_FLAGS BBLF_IS_ENTRY       =  0x2;
const BBL_FLAGS BBLF_IS_EXIT        =  0x4;
const BBL_FLAGS BBLF_IS_STAT_RET    =  0x8;
const BBL_FLAGS BBLF_IS_DYN_RET     = 0x10;
const BBL_FLAGS BBLF_DUMMY_CALL_RET = 0x20;

typedef struct _BBL_CTX _BBL_CTX, *BBL_CTX;
struct _BBL_CTX
{
    uint64_t exec_ct;
    /* preds and succs are arrays of indices into the edges arrays in the global context */
    uint32_t id, img_puid, pred_ct, pred_cap, succ_ct, succ_cap, ins_ct, *preds, *succs;
    INSTR inss;
    BBL_FLAGS flags;
};

typedef uint32_t RET_FLAGS;
const RET_FLAGS RETF_NONE = 0x0;
const RET_FLAGS RETF_CALL = 0x1;

typedef struct _RET_ADDR_CTX _RET_ADDR_CTX, *RET_ADDR_CTX;
struct _RET_ADDR_CTX
{
    ADDRINT addr;
    uint32_t bbl_id;
    RET_FLAGS flags;
};

typedef struct _IMG_CTX _IMG_CTX, *IMG_CTX;
struct _IMG_CTX
{
    char *name;
    uint32_t pin_id, puid;
    bool loaded;
};

typedef struct _QEXP_CTX _QEXP_CTX, *QEXP_CTX;
struct _QEXP_CTX
{
    std::ofstream *out, *log;
    uint64_t good_bbl_exec_ct, bad_bbl_exec_ct;
    uint32_t prev_bbl, bbl_ct, bbl_cap, edge_ct, edge_cap, ret_addr_ct, ret_addr_cap,
        img_ct, img_cap;
    BBL_CTX bbls;
    EDGE edges;
    RET_ADDR_CTX ret_addrs;
    IMG_CTX imgs;
    PIN_LOCK lock;
    pid_t parent_pid;
};

const UINT16 MAX_REG = 4096;
const uint32_t FIRST_REAL_BBL_IDX = 10;

static void add_edge( QEXP_CTX ctx, BBL_CTX from, BBL_CTX to, EDGE *edge_new )
{
    uint32_t new_edge_id = ctx->edge_ct;
    // printf( "AE1 %p  ", from->succs ); fflush( stdout );
    CHECK_CAP_REALLOC( from->succs, uint32_t*, from->succ_ct, from->succ_cap );
    // printf( "AE2 %p  ", to->preds ); fflush( stdout );
    CHECK_CAP_REALLOC( to->preds, uint32_t*, to->pred_ct, to->pred_cap );
    // printf( "AE3 %p  ", ctx->edges ); fflush( stdout );
    CHECK_CAP_REALLOC( ctx->edges, EDGE, ctx->edge_ct, ctx->edge_cap );
    EDGE edge_tmp = &ctx->edges[ new_edge_id ];
    from->succs[ from->succ_ct ] = new_edge_id;
    to->preds[ to->pred_ct ] = new_edge_id;
    edge_tmp->id = new_edge_id;
    edge_tmp->from = from->id;
    edge_tmp->to = to->id;
    ++from->succ_ct;
    ++to->pred_ct;
    ++ctx->edge_ct;
    if( edge_new )
    {
        *edge_new = edge_tmp;
    }
}

static void add_pred( QEXP_CTX ctx, uint32_t from, BBL_CTX to, EDGE *edge_new )
{
    // printf( "AP " ); fflush( stdout );
    add_edge( ctx, &ctx->bbls[ from ], to, edge_new );
}

static void add_succ( QEXP_CTX ctx, BBL_CTX from, uint32_t to, EDGE *edge_new )
{
    // printf( "AS " ); fflush( stdout );
    add_edge( ctx, from, &ctx->bbls[ to ], edge_new );
}

static void record_succ( uint32_t succ_id, QEXP_CTX ctx )
{
    uint32_t i;
    if( ctx->prev_bbl >= FIRST_REAL_BBL_IDX )
    {
        BBL_CTX bbl = &ctx->bbls[ ctx->prev_bbl ];
        for( i = 0; i < bbl->succ_ct; ++i )
        {
            if( ctx->edges[ bbl->succs[ i ] ].to == succ_id )
            {
                break;
            }
        }
        if( i == bbl->succ_ct )
        {
            // printf( "RS->AS %i ", succ_id ); fflush( stdout );
            add_succ( ctx, bbl, succ_id, NULL );
        }
#if 0
        if( i > 0 )
        {
            /* move-to-front style */
            _EDGE temp = bbl->succs[ i ];
            bbl->succs[ i ] = bbl->succs[ 0 ];
            bbl->succs[ 0 ] = temp;
        }
        ++bbl->succs[ 0 ].exec_ct;
#else
#endif
    }
}

static void on_bad_bbl_exec( void *v )
{
    QEXP_CTX ctx = (QEXP_CTX)v;
    ++ctx->bad_bbl_exec_ct;
}

static void on_good_bbl_exec( uint64_t bbl_id, void *v )
{
    QEXP_CTX ctx = (QEXP_CTX)v;
    ++ctx->good_bbl_exec_ct;
    if( ctx->prev_bbl )
    {
        BBL_CTX bbl = &ctx->bbls[ bbl_id ];
        uint32_t i;
        ++bbl->exec_ct;
        for( i = 0; i < bbl->pred_ct; ++i )
        {
            if( ctx->edges[ bbl->preds[ i ] ].from == ctx->prev_bbl )
            {
                break;
            }
        }
        if( i == bbl->pred_ct )
        {
            add_pred( ctx, ctx->prev_bbl, bbl, NULL );
        }
#if 0
        if( i > 0 )
        {
            /* move-to-front style */
            _EDGE temp = bbl->preds[ i ];
            bbl->preds[ i ] = bbl->preds[ 0 ];
            bbl->preds[ 0 ] = temp;
        }
        ++bbl->preds[ 0 ].exec_ct;
#else
        ++ctx->edges[ bbl->preds[ i ] ].exec_ct;
#endif
        record_succ( bbl_id, ctx );
    }
    ctx->prev_bbl = bbl_id;
}

void on_syscall_entry(
    THREADID threadIndex, CONTEXT *ctxt, SYSCALL_STANDARD std, VOID *v )
{
    QEXP_CTX ctx = (QEXP_CTX)v;
    if( ctx->prev_bbl )
    {
        record_succ( 1, ctx );
    }
    ctx->prev_bbl = 1;
}

static void add_bbl( QEXP_CTX ctx, BBL_CTX *bbl_ctx )
{
    CHECK_CAP_REALLOC( ctx->bbls, BBL_CTX, ctx->bbl_ct, ctx->bbl_cap );
    BBL_CTX temp = &ctx->bbls[ ctx->bbl_ct ];
    temp->id = ctx->bbl_ct;
    temp->pred_cap = 4;
    temp->preds =
        (uint32_t*)ff_calloc( temp->pred_cap, sizeof( temp->preds[0] ) );
    temp->succ_cap = 4;
    temp->succs =
        (uint32_t *)ff_calloc( temp->succ_cap, sizeof( temp->succs[0] ) );
    *bbl_ctx = temp;
    ++ctx->bbl_ct;
}

IMG_CTX find_img_by_pin_id( uint32_t pin_id, QEXP_CTX ctx )
{
    for( uint32_t i = 0; i < ctx->img_ct; ++i )
    {
        if( ctx->imgs[i].pin_id == pin_id )
            return &ctx->imgs[i];
    }
    return NULL;
}

static bool xed_operand_is_mem( xed_operand_enum_t name )
{
    return name == XED_OPERAND_MEM0 || name == XED_OPERAND_MEM1;
}

static void xed_stuff( INS ins, QEXP_CTX ctx ) VARIABLE_IS_NOT_USED;
static void xed_stuff( INS ins, QEXP_CTX ctx )
{
    const xed_decoded_inst_t* xedd = INS_XedDec( ins );
    const xed_inst_t* xedi = xed_decoded_inst_inst( xedd );

    for( uint32_t o = 0 ; o < xed_inst_noperands( xedi ); ++o )
    {
        const int op_str_len = 256;
        char op_str[ op_str_len ];
        const xed_operand_t* xop = xed_inst_operand( xedi, o );
        const xed_operand_enum_t operand_name = xed_operand_name( xop );
        xed_operand_print( xop, op_str, op_str_len );
        *ctx->log << "  xop " << o << " " << op_str;
        if( xed_operand_is_register( operand_name ) )
        {
            xed_reg_enum_t re = xed_decoded_inst_get_reg( xedd, operand_name );
            *ctx->log << " IS_REG! " << re << "(sp:" << XED_REG_SP << " e:"
                     << XED_REG_ESP << " r:" << XED_REG_RSP << ")";
        }
        else if( xed_operand_is_memory_addressing_register( operand_name ) )
        {
            xed_reg_enum_t re = xed_decoded_inst_get_reg( xedd, operand_name );
            *ctx->log << " IS_MEM_REG! " << re << "(sp:" << XED_REG_SP << " e:"
                     << XED_REG_ESP << " r:" << XED_REG_RSP << ")";
        }
        else if( xed_operand_is_mem( operand_name ) )
        {
            xed_reg_enum_t bre = xed_decoded_inst_get_base_reg(
                xedd, operand_name - XED_OPERAND_MEM0 );
            xed_reg_enum_t ire = xed_decoded_inst_get_index_reg(
                xedd, operand_name - XED_OPERAND_MEM0 );
            xed_reg_enum_t sre = xed_decoded_inst_get_seg_reg(
                xedd, operand_name - XED_OPERAND_MEM0 );
            *ctx->log << " IS_MEM! B:" << bre << " I:" << ire << " S:" << sre;
        }
        *ctx->log << std::endl;
    }
}

/*
 * Main instrumentation entry point.
 */
void bbl_translate( QEXP_CTX ctx, IMG_CTX img, BBL bbl, bool first );
VOID on_trace_translation( TRACE trace, void *v )
{
    bool first = true;
    QEXP_CTX ctx = (QEXP_CTX)v;
    const RTN rtn = TRACE_Rtn( trace );
    if( !RTN_Valid( rtn ) )
    {
        // IMG img = IMG_FindByAddress( TRACE_Address( trace ) );
        for( BBL bbl = TRACE_BblHead( trace );
             BBL_Valid( bbl );
             bbl = BBL_Next( bbl ) )
        {
            BBL_InsertCall(
                bbl, IPOINT_ANYWHERE, AFUNPTR( on_bad_bbl_exec ),
                IARG_PTR, ctx, IARG_END);
        }
        return;
    }
    const SEC sec = RTN_Sec( rtn );
    ASSERTX( SEC_Valid( sec ) );
    const IMG pin_img = SEC_Img( sec );
    ASSERTX( IMG_Valid( pin_img ) );
    IMG_CTX img = find_img_by_pin_id( IMG_Id( pin_img ), ctx );
    ASSERTX( img && img->loaded );

    // *ctx->log << "Instrumenting BBL! pid=" << PIN_GetPid() << std::endl;

    for( BBL bbl = TRACE_BblHead( trace );
         BBL_Valid( bbl );
         bbl = BBL_Next( bbl ) )
    {
        bbl_translate( ctx, img, bbl, first );
        first = false;
    }
}

void instruction_translate( QEXP_CTX ctx, BBL_CTX bbl_ctx, INS ins, uint32_t &i );
void bbl_translate( QEXP_CTX ctx, IMG_CTX img, BBL bbl, bool first )
{
    BBL_CTX bbl_ctx;
    INS tail = BBL_InsTail( bbl );
    // *ctx->out << "  Doing a BBL " << INS_Address( BBL_InsHead( bbl ) ) << std::endl;

    if( first )
    {
        CHECK_CAP_REALLOC(
            ctx->ret_addrs, RET_ADDR_CTX, ctx->ret_addr_ct, ctx->ret_addr_cap );
        RET_ADDR_CTX ret_addr = &ctx->ret_addrs[ ctx->ret_addr_ct ];
        ret_addr->addr = INS_Address( BBL_InsHead( bbl ) );
        ret_addr->bbl_id = ctx->bbl_ct;
        ++ctx->ret_addr_ct;
    }

    add_bbl( ctx, &bbl_ctx );

    bbl_ctx->img_puid = img->puid;
    if( INS_IsCall( tail ) || INS_IsProcedureCall( tail ) || INS_IsSyscall( tail ) )
    {
        bbl_ctx->flags |= BBLF_IS_CALL;
        CHECK_CAP_REALLOC(
            ctx->ret_addrs, RET_ADDR_CTX, ctx->ret_addr_ct, ctx->ret_addr_cap );
        RET_ADDR_CTX ret_addr = &ctx->ret_addrs[ ctx->ret_addr_ct ];
        ret_addr->addr = INS_NextAddress( tail );
        ret_addr->bbl_id = bbl_ctx->id;
        ret_addr->flags |= RETF_CALL;
        ++ctx->ret_addr_ct;
    }

    if( INS_IsRet( tail ) || INS_IsSysret( tail ) )
    {
        bbl_ctx->flags |= BBLF_IS_EXIT;
    }

    for( INS ins = BBL_InsHead( bbl );
         INS_Valid( ins );
         ins = INS_Next( ins ) )
    {
        ++bbl_ctx->ins_ct;
    }

    bbl_ctx->inss = (INSTR)ff_calloc(
        bbl_ctx->ins_ct, sizeof( bbl_ctx->inss[0] ) );

    uint32_t i = 0;
    for( INS ins = BBL_InsHead( bbl );
         INS_Valid( ins );
         ins = INS_Next( ins ) )
    {
        instruction_translate( ctx, bbl_ctx, ins, i );
    }
    BBL_InsertCall(
        bbl, IPOINT_ANYWHERE, AFUNPTR( on_good_bbl_exec ),
        IARG_UINT32, bbl_ctx->id,
        IARG_PTR,    ctx,
        IARG_END);
}

void instruction_translate( QEXP_CTX ctx, BBL_CTX bbl_ctx, INS ins, uint32_t &i )
{
    string disas = INS_Disassemble( ins );
    size_t disas_len = disas.length();
    INSTR instr = &bbl_ctx->inss[i];
    _ANNOT_REG reads_temp[100], writes_temp[100];
    uint8_t reads_temp_ct = 0, writes_temp_ct = 0;
    bool any_mystery_operands = false;

    REG const reg_inv = REG_INVALID();
    instr->disas = (char *)ff_malloc( disas_len + 1 );
    disas.copy( instr->disas, disas_len );
    instr->disas[ disas_len ] = 0;
    instr->opc = INS_Opcode( ins );
    if( INS_IsMemoryRead( ins ) )
        instr->flags |= INSTRF_MEM_READ;
    if( INS_IsMemoryWrite( ins ) )
        instr->flags |= INSTRF_MEM_WRITE;
    if( INS_HasFallThrough( ins ) )
        instr->flags |= INSTRF_HAS_FALL;

    // *ctx->log << INS_Disassemble( ins ) << " " << std::endl;

    for( uint32_t o = 0; o < INS_OperandCount( ins ); ++o )
    {
        if( INS_OperandIsReg( ins, o ) )
        {
            AREG_FLAGS f = REGF_NONE;
            if( INS_OperandIsImplicit( ins, o ) )
                f |= REGF_IMPLICIT;
            if( INS_OperandReadAndWritten( ins, o ) )
                f |= REGF_READ_WRITE;
            if( INS_OperandRead( ins, o ) )
            {
                reads_temp[ reads_temp_ct ].r = INS_OperandReg( ins, o );
                reads_temp[ reads_temp_ct ].flags = f;
                ++reads_temp_ct;
            }
            if( INS_OperandWritten( ins, o ) )
            {
                writes_temp[ writes_temp_ct ].r = INS_OperandReg( ins, o );
                writes_temp[ writes_temp_ct ].flags = f;
                ++writes_temp_ct;
            }
        }
        else if( INS_OperandIsMemory( ins, o )
                 || INS_OperandIsAddressGenerator( ins, o ) )
        {
            AREG_FLAGS f =
                INS_OperandIsImplicit( ins, o ) ? REGF_IMPLICIT : REGF_NONE;
            if( reg_inv != INS_OperandMemoryBaseReg( ins, o ) )
            {
                reads_temp[ reads_temp_ct ].r = INS_OperandMemoryBaseReg( ins, o );
                reads_temp[ reads_temp_ct ].flags = f;
                ++reads_temp_ct;
            }
            if( reg_inv != INS_OperandMemoryIndexReg( ins, o ) )
            {
                reads_temp[ reads_temp_ct ].r = INS_OperandMemoryIndexReg( ins, o );
                reads_temp[ reads_temp_ct ].flags = f;
                ++reads_temp_ct;
            }
            if( reg_inv != INS_OperandMemorySegmentReg( ins, o ) )
            {
                reads_temp[ reads_temp_ct ].r = INS_OperandMemorySegmentReg( ins, o );
                reads_temp[ reads_temp_ct ].flags = f;
                ++reads_temp_ct;
            }
        }
    }

    instr->r_ct = INS_MaxNumRRegs( ins );
    instr->w_ct = INS_MaxNumWRegs( ins );
    instr->reads  = (ANNOT_REG)ff_calloc(
        instr->r_ct, sizeof( instr->reads[0] ) );
    instr->writes = (ANNOT_REG)ff_calloc(
        instr->w_ct, sizeof( instr->writes[0] ) );
    // *ctx->log << " rs:";
    for( uint32_t r = 0; r < instr->r_ct; ++r )
    {
        uint8_t orig_temp_ct = reads_temp_ct;
        instr->reads[ r ].r = INS_RegR( ins, r );
        // *ctx->log << " " << REG_StringShort( instr->reads[ r ].r );
        for( uint8_t rt = 0; rt < reads_temp_ct; ++rt )
        {
            if( instr->reads[ r ].r == reads_temp[ rt ].r )
            {
                instr->reads[ r ].flags = reads_temp[ rt ].flags;
                for( uint8_t rt2 = rt + 1; rt2 < reads_temp_ct; ++rt2 )
                    reads_temp[ rt2 - 1 ] = reads_temp[ rt2 ];
                --reads_temp_ct;
                break;
            }
        }
        // *ctx->log << " f" << instr->reads[ r ].flags;
        if( orig_temp_ct == reads_temp_ct )
        {
            instr->reads[ r ].flags = REGF_IMPLICIT;
            *ctx->log << " " << REG_StringShort( instr->reads[ r ].r );
            *ctx->log << instr->reads[ r ].flags;
            any_mystery_operands = true;
        }
    }
    // *ctx->log << std::endl;
    if( reads_temp_ct > 0 )
    {
        xed_stuff( ins, ctx );
        instr->flags |= INSTRF_TOO_MANY_READ_OPERANDS;
    }
    //*ctx->log << " ws:";
    for( uint32_t w = 0; w < instr->w_ct; ++w )
    {
        uint8_t orig_temp_ct = writes_temp_ct;
        instr->writes[ w ].r = INS_RegW( ins, w );
        //*ctx->log << " " << REG_StringShort( instr->writes[ w ].r );
        for( uint8_t wt = 0; wt < writes_temp_ct; ++wt )
        {
            if( instr->writes[ w ].r == writes_temp[ wt ].r )
            {
                instr->writes[ w ].flags = writes_temp[ wt ].flags;
                for( uint8_t wt2 = wt + 1; wt2 < writes_temp_ct; ++wt2 )
                    writes_temp[ wt2 - 1 ] = writes_temp[ wt2 ];
                --writes_temp_ct;
                break;
            }
        }
        // *ctx->log << instr->writes[ w ].flags;
        if( orig_temp_ct == writes_temp_ct )
        {
            instr->writes[ w ].flags = REGF_IMPLICIT;
            *ctx->log << " " << REG_StringShort( instr->writes[ w ].r );
            *ctx->log << instr->writes[ w ].flags;
            any_mystery_operands = true;
        }
    }
    if( writes_temp_ct > 0 )
    {
        xed_stuff( ins, ctx );
        instr->flags |= INSTRF_TOO_MANY_WRITE_OPERANDS;
    }
    if( any_mystery_operands )
    {
        *ctx->log << "   " << INS_Disassemble( ins ) << " " << std::endl;
        xed_stuff( ins, ctx );
        *ctx->log << std::endl;
    }
    // *ctx->log << std::endl;
    ++i;
}

int bbl_compare( const void *p1, const void *p2 )
{
    BBL_CTX b1 = (BBL_CTX)p1, b2 = (BBL_CTX)p2;
    int diff = b2->ins_ct * b2->exec_ct - b1->ins_ct * b1->exec_ct;
    if( diff == 0 )
    {
        if( b1->exec_ct > 0 )
        {
            diff = b2->ins_ct - b1->ins_ct;
        }
        if( diff == 0 )
        {
            diff = b1->id - b2->id;
        }
    }
    return diff;
}

int bbl_id_compare( const void *p1, const void *p2 )
{
    BBL_CTX b1 = (BBL_CTX)p1, b2 = (BBL_CTX)p2;
    return b2->id > b1->id ? -1 : ( b2->id < b1->id ? 1 : 0 );
}

int ret_addr_compare( const void *p1, const void *p2 )
{
    RET_ADDR_CTX b1 = (RET_ADDR_CTX)p1, b2 = (RET_ADDR_CTX)p2;
    int diff = b2->addr > b1->addr ? 1 : ( b2->addr < b1->addr ? -1 : 0 );
    if( diff == 0 )
    {
        diff = b2->flags - b1->flags;
        if( diff == 0 )
        {
            diff = b2->bbl_id - b1->bbl_id;
        }
    }
    return diff;
}

static bool bbl_matters( BBL_CTX bbl )
{
    return bbl->exec_ct > 0
        || bbl->flags & BBLF_DUMMY_CALL_RET;
}

void on_process_termination( int, void *v )
{
    uint64_t bct = 0, ict = 0, rct = 0;
    QEXP_CTX ctx = (QEXP_CTX)v;
    uint32_t i;
    bool first = true;
    ADDRINT curr_addr = 0;
    BBL_CTX call_link_ctx = NULL;

    GetLock( &ctx->lock, 0 );
    *ctx->log << "  Reached Fini!  :)BBL " << ctx->good_bbl_exec_ct << ",  :(BBL "
              << ctx->bad_bbl_exec_ct << " pid: " << PIN_GetPid() << std::endl;
    ReleaseLock( &ctx->lock );

    qsort( ctx->ret_addrs, ctx->ret_addr_ct, sizeof( ctx->ret_addrs[0] ), ret_addr_compare );

    for( i = 0; i < ctx->ret_addr_ct; ++i )
    {
        RET_ADDR_CTX ret_addr = &ctx->ret_addrs[ i ];
        BBL_CTX other_ctx = &ctx->bbls[ ret_addr->bbl_id ];
        if( ret_addr->flags & RETF_CALL )
        {
            if( curr_addr != ret_addr->addr )
            {
                add_bbl( ctx, &call_link_ctx );
                /* bbls might have moved (because of realloc); recalc other_ctx */
                other_ctx = &ctx->bbls[ ret_addr->bbl_id ];
                call_link_ctx->flags |= BBLF_DUMMY_CALL_RET;
                // *ctx->out << std::endl;
            }
            add_edge( ctx, other_ctx, call_link_ctx, NULL );
            curr_addr = ret_addr->addr;
        }
        if( curr_addr == ret_addr->addr )
        {
            ASSERTX( call_link_ctx != NULL );
            if( !( ret_addr->flags & RETF_CALL ) )
            {
                add_edge( ctx, call_link_ctx, other_ctx, NULL );
                other_ctx->flags |= BBLF_IS_STAT_RET;
            }
            //*ctx->out << ret_addr->flags << " " << ret_addr->addr << " "
            //          << ret_addr->bbl_id << " " << std::endl;
        }
        else
        {
            call_link_ctx = NULL;
        }
    }

    // uint64_t blah = 0;
    for( i = FIRST_REAL_BBL_IDX; i < ctx->bbl_ct; ++i )
    {
        BBL_FLAGS flags = BBLF_NONE;
        size_t j;
        BBL_CTX bbl = &ctx->bbls[ i ];
        bool is_d_ret, is_s_ret;
        for( j = 0; j < bbl->pred_ct; ++j )
        {
            uint32_t pred_id = ctx->edges[ bbl->preds[ j ] ].from;
            flags |= ctx->bbls[ pred_id ].flags;
        }
        if( flags & BBLF_IS_EXIT )
        {
            bbl->flags |= BBLF_IS_DYN_RET;
        }
        if( flags & BBLF_IS_CALL )
        {
            bbl->flags |= BBLF_IS_ENTRY;
        }
        // *ctx->out << "  " << bbl->id << "  " << bbl->flags << std::endl;
        is_d_ret = bbl->flags & BBLF_IS_DYN_RET;
        is_s_ret = bbl->flags & BBLF_IS_STAT_RET;
        if( ( is_d_ret || is_s_ret ) && !( is_d_ret && is_s_ret ) )
        {
            //*ctx->out << "  WEIRD " << bbl->exec_ct << std::endl;
            // blah += bbl->exec_ct * bbl->ins_ct;
        }
        if( is_d_ret && !is_s_ret )
        {
#if 0
            if( call_link_ctx == NULL )
            {
                add_bbl( ctx, &call_link_ctx );
                call_link_ctx->flags |= BBLF_DUMMY_CALL_RET;
                *ctx->out << std::endl;
            }
#endif
        }
    }
    // *ctx->out << "    Total instrs executed by weird BBLs " << blah << std::endl;

    *ctx->log << "Before first sort" << std::endl;
    for( i = 0; i < FIRST_REAL_BBL_IDX * 2; ++i )
    {
        BBL_CTX bbl = &ctx->bbls[i];
        *ctx->log << i << "  " << bbl->id << std::endl;
    }

    uint32_t actual_image_count = 0, actual_bbl_count = 0, actual_edge_count = 0,
        max_image_id = 0, max_bbl_id = 0, max_edge_id = 0;

    for( i = 0; i < ctx->img_ct; ++i )
    {
        ++actual_image_count;
        max_image_id = MAX( max_image_id, ctx->imgs[ i ].puid );
    }
    for( i = FIRST_REAL_BBL_IDX; i < ctx->bbl_ct; ++i )
    {
        BBL_CTX bbl = &ctx->bbls[i];
        if( bbl_matters( bbl ) )
        {
            ++actual_bbl_count;
            max_bbl_id = MAX( max_bbl_id, bbl->id );
        }
    }
    for( i = 0; i < ctx->edge_ct; ++i )
    {
        EDGE edge = &ctx->edges[i];
        if( bbl_matters( &ctx->bbls[ edge->from ] )
            && bbl_matters( &ctx->bbls[ edge->to ] ) )
        {
            ++actual_edge_count;
            max_edge_id = MAX( max_edge_id, edge->id );
        }
    }

    /* !!! This sort invalidates the invariant that BBL_id is at
     * ctx->bbls[id].  Be very careful with indexing into the bbls array
     * from here on. */
    qsort( &ctx->bbls[ FIRST_REAL_BBL_IDX ],
           ctx->bbl_ct - FIRST_REAL_BBL_IDX,
           sizeof( ctx->bbls[0] ),
           bbl_compare );

    *ctx->log << "After first sort" << std::endl;
    for( i = 0; i < FIRST_REAL_BBL_IDX * 2; ++i )
    {
        BBL_CTX bbl = &ctx->bbls[i];
        *ctx->log << i << "  " << bbl->id << std::endl;
    }

#define write_prop( name, prop, quotes, comma )       \
    do { \
        if( quotes ) *ctx->out << "\""#name"\":\"" << ( prop ) << ( comma ? "\"," : "\"" ); \
        else         *ctx->out << "\""#name"\":"   << ( prop ) << ( comma ? "," : "" ); \
    } while( false )

    *ctx->out << "{ ";
    write_prop( ni,  actual_image_count, false, true );
    write_prop( mii, max_image_id,       false, true );
    write_prop( nb,  actual_bbl_count,   false, true );
    write_prop( mbi, max_bbl_id,         false, true );
    write_prop( ne,  actual_edge_count,  false, true );
    write_prop( mei, max_edge_id,        false, true );
    *ctx->out << std::endl << "\"images\": [" << std::endl;
    for( i = 0; i < ctx->img_ct; ++i )
    {
        IMG_CTX img = &ctx->imgs[ i ];
        if( !first )
        {
            *ctx->out << "," << std::endl;
        }
        first = false;
        *ctx->out << "  { ";
        write_prop( name,   img->name,   true,  true );
        write_prop( pin_id, img->pin_id, false, true );
        write_prop( puid,   img->puid,   false, false );
        *ctx->out << " }";
    }
    first = true;
    *ctx->out << std::endl << "], \"bbls\": [" << std::endl;
    for( i = FIRST_REAL_BBL_IDX; i < ctx->bbl_ct; ++i )
    {
        BBL_CTX bbl = &ctx->bbls[i];
        ++bct;
        if( !bbl_matters( bbl ) )
            continue;
        if( !first )
            *ctx->out << "," << std::endl;
        first = false;
        *ctx->out << "  {" << std::endl << "    ";
        write_prop( id,     bbl->id,       false, true );
        write_prop( exec,   bbl->exec_ct,  false, true );
        write_prop( insN,   bbl->ins_ct,   false, true );
        write_prop( bfl,    bbl->flags,    false, true );
        /* These counts aren't exctly right because of "doesn't matter" */
        write_prop( np,     bbl->pred_ct,  false, true );
        write_prop( ns,     bbl->succ_ct,  false, true );
        write_prop( img_id, bbl->img_puid, false, false );
        *ctx->out << ",\"inss\":" << std::endl << "    [";
        if( bbl->ins_ct > 0 )
        {
            for( uint32_t i = 0; i < bbl->ins_ct; ++i )
            {
                INSTR ins = &bbl->inss[ i ];
                ++ict;
                if( i > 0 )
                {
                    *ctx->out << ",";
                }

                *ctx->out << std::endl << "      { ";
                write_prop( ifl, ins->flags, false, true );
                write_prop( cd, OPCODE_StringShort( ins->opc ), true, true );
                write_prop( nr, ((int)ins->r_ct), false, true );
                write_prop( nw, ((int)ins->w_ct), false, true );
                *ctx->out << "\"r\":[";
                for( uint32_t r = 0; r < ins->r_ct; ++r )
                {
                    const REG reg = ins->reads[r].r;
                    ++rct;
                    if( r > 0 )
                    {
                        *ctx->out << ",";
                    }
                    /* XXX: Figure out if FullRegName is the right thing to do here */
                    *ctx->out << std::endl << "          {";
                    write_prop( r, REG_StringShort( reg ), true, true );
                    write_prop( rf, REG_StringShort( REG_FullRegName( reg ) ), true, true );
                    write_prop( fl, ins->reads[r].flags, false, false );
                    *ctx->out << "}";
                }
                *ctx->out << std::endl << "        ], \"w\":[";
                for( uint32_t r = 0; r < ins->w_ct; ++r )
                {
                    const REG reg = ins->writes[r].r;
                    ++rct;
                    if( r > 0 )
                    {
                        *ctx->out << ",";
                    }
                    *ctx->out << std::endl << "          {";
                    write_prop( r, REG_StringShort( reg ), true, true );
                    write_prop( rf, REG_StringShort( REG_FullRegName( reg ) ), true, true );
                    write_prop( fl, ins->reads[r].flags, false, false );
                    *ctx->out << "}";
                }
                *ctx->out << std::endl << "        ], ";
                write_prop( ds, ins->disas, true, false );
                *ctx->out << " }";
            }
        }
        *ctx->out << std::endl << "    ]";
        *ctx->out << std::endl << "  }";
    }
    *ctx->out << std::endl << "], \"edges\": [" << std::endl;

    qsort( &ctx->bbls[ FIRST_REAL_BBL_IDX ],
           ctx->bbl_ct - FIRST_REAL_BBL_IDX,
           sizeof( ctx->bbls[0] ),
           bbl_id_compare );

    *ctx->log << "After second sort" << std::endl;
    for( i = 0; i < FIRST_REAL_BBL_IDX * 2; ++i )
    {
        BBL_CTX bbl = &ctx->bbls[i];
        *ctx->log << i << "  " << bbl->id << std::endl;
    }

    first = true;
    for( i = 0; i < ctx->edge_ct; ++i )
    {
        EDGE edge = &ctx->edges[i];

        if( !( bbl_matters( &ctx->bbls[ edge->from ] )
               && bbl_matters( &ctx->bbls[ edge->to ] ) ) )
        {
            continue;
        }
        if( !first )
        {
            *ctx->out << "," << std::endl;
        }
        first = false;
        if( edge->from != ctx->bbls[ edge->from ].id
            || edge->to != ctx->bbls[ edge->to ].id )
        {
            /* error? */
        }
        *ctx->out << "  { ";
        write_prop( id,   edge->id,      false, true );
        write_prop( ect,  edge->exec_ct, false, true );
        write_prop( from, edge->from,    false, true );
        write_prop( to,   edge->to,      false, false );
        *ctx->out << " }";
    }
    *ctx->out << std::endl << "]" << std::endl << "}" << std::endl;
    // *ctx->out << std::endl << "b:" << bct << "  e:" << ect << "  i:" << ict << "  r:" << rct << std::endl;

    ctx->out->close();
}

void reset_counts( QEXP_CTX ctx )
{
    uint32_t i;
    ctx->good_bbl_exec_ct = 0;
    ctx->bad_bbl_exec_ct = 0;
    for( i = 0; i < ctx->bbl_ct; ++i )
    {
        ctx->bbls[i].exec_ct = 0;
    }
    for( i = 0; i < ctx->edge_ct; ++i )
    {
        ctx->edges[i].exec_ct = 0;
    }
}

void open_output_file( std::ofstream **o )
{
    int pid = PIN_GetPid(), i = 1;
    const int sz_lim = 1000;
    char fname[ sz_lim ];
    if( !o )
        PIN_ExitProcess( EBADF );
    while( true )
    {
        int sz = snprintf( fname, sz_lim, "%s_%i_%i.%s", KnobOutputFileBase.Value().c_str(),
                           pid, i, KnobOutputFileSuffix.Value().c_str() );
        if( sz >= sz_lim )
        {
            PIN_ExitProcess( ENOSPC );
        }
        ifstream my_file( fname );
        if( my_file.good() )
        {
            my_file.close();
            ++i;
        }
        else
        {
            break;
        }
    }
    *o = new std::ofstream( fname );
}

void on_pre_fork( THREADID threadid, const CONTEXT* ctxt, void *arg )
{
    QEXP_CTX ctx = (QEXP_CTX)arg;
    GetLock( &ctx->lock, threadid+1 );
    *ctx->log << "TOOL: Before fork." << endl;
    ReleaseLock( &ctx->lock );
    ctx->parent_pid = PIN_GetPid();
}

void on_post_fork_child(
    THREADID threadid, const CONTEXT* ctxt, void *arg )
{
    QEXP_CTX ctx = (QEXP_CTX)arg;
    reset_counts( ctx );
    GetLock( &ctx->lock, threadid+1 );
    *ctx->log << "TOOL: After fork in child. " << getpid() << endl; ctx->log->flush();
    ReleaseLock( &ctx->lock );

    //ctx->out->close();
    open_output_file( &ctx->out );

    if( ( PIN_GetPid() == ctx->parent_pid ) || ( getppid() != ctx->parent_pid ) )
    {
        cerr << "PIN_GetPid() fails in child process" << endl;
        PIN_ExitProcess( ECHILD );
    }
}

void init_ctx( QEXP_CTX ctx )
{
    BBL_CTX bbl_special;
    size_t i;
    open_output_file( &ctx->out );
    ctx->log = new std::ofstream( KnobOutputLogName.Value().c_str() );
    ctx->bbl_cap = FIRST_REAL_BBL_IDX * 2;
    ctx->bbls = (BBL_CTX)ff_calloc( ctx->bbl_cap, sizeof( ctx->bbls[0] ) );
    ctx->edge_cap = 10;
    ctx->edges = (EDGE)ff_calloc( ctx->edge_cap, sizeof( ctx->edges[0] ) );
    ctx->ret_addr_cap = 20;
    ctx->ret_addrs = (RET_ADDR_CTX)ff_calloc( ctx->ret_addr_cap, sizeof( ctx->ret_addrs[0] ) );
    ctx->img_cap = 10;
    ctx->imgs = (IMG_CTX)ff_calloc( ctx->img_cap, sizeof( ctx->imgs[0] ) );
    for( i = 0; i < FIRST_REAL_BBL_IDX; ++i )
    {
        ctx->bbls[ i ].id = i;
    }
    ctx->bbl_ct = 1;
    add_bbl( ctx, &bbl_special );
    // bbl_special->... XXX: syscall
    ctx->bbl_ct = FIRST_REAL_BBL_IDX;
    InitLock( &ctx->lock );
}

BOOL on_child_exec( CHILD_PROCESS childProcess, void *a )
{
    on_process_termination( 0, a );
    return true;
}

void on_image_load( IMG pin_img, void *arg )
{
    QEXP_CTX ctx = (QEXP_CTX)arg;
    IMG_CTX img = find_img_by_pin_id( IMG_Id( pin_img ), ctx );
    if( img )
    {
        if( img->loaded )
        {
            std::cerr << "Trying to load already loaded image!!" << std::endl;
            PIN_ExitProcess( EMLINK );
        }
    }
    else
    {
        const string n = IMG_Name( pin_img );
        CHECK_CAP_REALLOC( ctx->imgs, IMG_CTX, ctx->img_ct, ctx->img_cap );
        img = &ctx->imgs[ ctx->img_ct ];
        img->name = (char *)ff_malloc( n.size() + 1 );
        strcpy( img->name, n.c_str() );
        img->pin_id = IMG_Id( pin_img );
        img->puid = ctx->img_ct;
        ++ctx->img_ct;
    }
    img->loaded = true;
}

void on_image_unload( IMG pin_img, void *arg )
{
    QEXP_CTX ctx = (QEXP_CTX)arg;
    IMG_CTX img = find_img_by_pin_id( IMG_Id( pin_img ), ctx );
    if( !img || !img->loaded )
    {
        std::cerr << "Trying to unload an image that doesn't exist or isn't loaded!!"
                  << std::endl;
        PIN_ExitProcess( EMLINK );
    }
    img->loaded = false;
}

int main( int argc, CHAR *argv[] )
{
    QEXP_CTX ctx;
    int i;
    PIN_InitSymbols();

    if( PIN_Init( argc, argv ) )
    {
        std::cerr << "Fail!!! " << argc << std::endl;
        for( i = 0; i < argc; ++i )
        {
            std::cerr << "  i[" << i << "] " << argv[i] << std::endl;
        }
        return Usage();
    }
    ctx = (QEXP_CTX)ff_malloc( sizeof( *ctx ) );
    init_ctx( ctx );
    *ctx->log << "Ben's miraculous pintool starts! " << argc
              << " probe?" << PIN_IsProbeMode() << std::endl;
    for( i = 0; i < argc; ++i )
    {
        *ctx->log << "  i[" << i << "] " << argv[i] << std::endl;
    }
    TRACE_AddInstrumentFunction               ( on_trace_translation,   ctx );
    PIN_AddSyscallEntryFunction               ( on_syscall_entry,       ctx );
    PIN_AddForkFunction( FPOINT_BEFORE,         on_pre_fork,            ctx );
    PIN_AddForkFunction( FPOINT_AFTER_IN_CHILD, on_post_fork_child,     ctx );
    PIN_AddFollowChildProcessFunction         ( on_child_exec,          ctx );
    IMG_AddInstrumentFunction                 ( on_image_load,          ctx );
    IMG_AddUnloadFunction                     ( on_image_unload,        ctx );
    PIN_AddFiniFunction                       ( on_process_termination, ctx );

    PIN_StartProgram();

    return 0;
}
